#!/bin/bash

# AIQuant 数据库初始化脚本
# 用于初始化AIQuant量化投研系统的数据库

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 默认配置
DEFAULT_DB_HOST="localhost"
DEFAULT_DB_PORT="5432"
DEFAULT_DB_NAME="aiquant"
DEFAULT_DB_USER="aiquant"
DEFAULT_DB_PASSWORD="aiquant123"

# 读取配置
read_config() {
    # 从环境变量读取配置
    DB_HOST=${DATABASE_HOST:-$DEFAULT_DB_HOST}
    DB_PORT=${DATABASE_PORT:-$DEFAULT_DB_PORT}
    DB_NAME=${DATABASE_NAME:-$DEFAULT_DB_NAME}
    DB_USER=${DATABASE_USER:-$DEFAULT_DB_USER}
    DB_PASSWORD=${DATABASE_PASSWORD:-$DEFAULT_DB_PASSWORD}
    
    # 从.env文件读取配置（如果存在）
    if [ -f ".env" ]; then
        source .env
        DB_HOST=${DATABASE_HOST:-$DB_HOST}
        DB_PORT=${DATABASE_PORT:-$DB_PORT}
        DB_NAME=${DATABASE_NAME:-$DB_NAME}
        DB_USER=${DATABASE_USER:-$DB_USER}
        DB_PASSWORD=${DATABASE_PASSWORD:-$DB_PASSWORD}
    fi
}

# 检查PostgreSQL连接
check_postgres() {
    log_info "检查PostgreSQL连接..."
    
    if command -v psql &> /dev/null; then
        if PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d postgres -c '\q' 2>/dev/null; then
            log_success "PostgreSQL连接成功"
            return 0
        else
            log_error "PostgreSQL连接失败"
            return 1
        fi
    else
        log_error "psql命令未找到，请安装PostgreSQL客户端"
        return 1
    fi
}

# 创建数据库
create_database() {
    log_info "创建数据库 $DB_NAME..."
    
    # 检查数据库是否已存在
    if PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d postgres -lqt | cut -d \| -f 1 | grep -qw $DB_NAME; then
        log_warning "数据库 $DB_NAME 已存在"
        return 0
    fi
    
    # 创建数据库
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d postgres -c "CREATE DATABASE $DB_NAME;"
    
    if [ $? -eq 0 ]; then
        log_success "数据库 $DB_NAME 创建成功"
    else
        log_error "数据库 $DB_NAME 创建失败"
        return 1
    fi
}

# 安装扩展
install_extensions() {
    log_info "安装数据库扩展..."
    
    # 安装TimescaleDB扩展（如果可用）
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -c "CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;" 2>/dev/null || {
        log_warning "TimescaleDB扩展不可用，将使用标准PostgreSQL"
    }
    
    # 安装其他有用的扩展
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -c "CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";" 2>/dev/null || true
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -c "CREATE EXTENSION IF NOT EXISTS \"pg_stat_statements\";" 2>/dev/null || true
    
    log_success "数据库扩展安装完成"
}

# 创建表结构
create_tables() {
    log_info "创建数据表结构..."
    
    # 创建SQL脚本
    cat > /tmp/create_tables.sql << 'EOF'
-- 用户表
CREATE TABLE IF NOT EXISTS users (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    username VARCHAR(50) UNIQUE NOT NULL,
    email VARCHAR(100) UNIQUE NOT NULL,
    password_hash VARCHAR(255) NOT NULL,
    is_active BOOLEAN DEFAULT true,
    is_admin BOOLEAN DEFAULT false,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 策略表
CREATE TABLE IF NOT EXISTS strategies (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    name VARCHAR(100) NOT NULL,
    description TEXT,
    strategy_type VARCHAR(50) NOT NULL,
    parameters JSONB,
    is_active BOOLEAN DEFAULT true,
    created_by UUID REFERENCES users(id),
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 市场数据表
CREATE TABLE IF NOT EXISTS market_data (
    id BIGSERIAL PRIMARY KEY,
    symbol VARCHAR(20) NOT NULL,
    timestamp TIMESTAMP WITH TIME ZONE NOT NULL,
    open_price DECIMAL(15,6),
    high_price DECIMAL(15,6),
    low_price DECIMAL(15,6),
    close_price DECIMAL(15,6),
    volume BIGINT,
    amount DECIMAL(20,6),
    data_source VARCHAR(50),
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 创建时序表（如果TimescaleDB可用）
SELECT create_hypertable('market_data', 'timestamp', if_not_exists => TRUE);

-- 交易订单表
CREATE TABLE IF NOT EXISTS orders (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    strategy_id UUID REFERENCES strategies(id),
    symbol VARCHAR(20) NOT NULL,
    order_type VARCHAR(20) NOT NULL,
    side VARCHAR(10) NOT NULL,
    quantity DECIMAL(15,6) NOT NULL,
    price DECIMAL(15,6),
    status VARCHAR(20) DEFAULT 'pending',
    filled_quantity DECIMAL(15,6) DEFAULT 0,
    average_price DECIMAL(15,6),
    commission DECIMAL(15,6) DEFAULT 0,
    order_time TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    filled_time TIMESTAMP WITH TIME ZONE,
    cancelled_time TIMESTAMP WITH TIME ZONE,
    external_order_id VARCHAR(100),
    broker VARCHAR(50),
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 持仓表
CREATE TABLE IF NOT EXISTS positions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    strategy_id UUID REFERENCES strategies(id),
    symbol VARCHAR(20) NOT NULL,
    quantity DECIMAL(15,6) NOT NULL,
    average_cost DECIMAL(15,6) NOT NULL,
    current_price DECIMAL(15,6),
    unrealized_pnl DECIMAL(15,6),
    realized_pnl DECIMAL(15,6) DEFAULT 0,
    last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    UNIQUE(strategy_id, symbol)
);

-- 交易记录表
CREATE TABLE IF NOT EXISTS trades (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    order_id UUID REFERENCES orders(id),
    strategy_id UUID REFERENCES strategies(id),
    symbol VARCHAR(20) NOT NULL,
    side VARCHAR(10) NOT NULL,
    quantity DECIMAL(15,6) NOT NULL,
    price DECIMAL(15,6) NOT NULL,
    commission DECIMAL(15,6) DEFAULT 0,
    trade_time TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    external_trade_id VARCHAR(100),
    broker VARCHAR(50),
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 风险指标表
CREATE TABLE IF NOT EXISTS risk_metrics (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    strategy_id UUID REFERENCES strategies(id),
    metric_name VARCHAR(50) NOT NULL,
    metric_value DECIMAL(15,6) NOT NULL,
    threshold_value DECIMAL(15,6),
    status VARCHAR(20) DEFAULT 'normal',
    calculated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 系统日志表
CREATE TABLE IF NOT EXISTS system_logs (
    id BIGSERIAL PRIMARY KEY,
    level VARCHAR(20) NOT NULL,
    message TEXT NOT NULL,
    module VARCHAR(100),
    function_name VARCHAR(100),
    line_number INTEGER,
    extra_data JSONB,
    timestamp TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 创建时序表（如果TimescaleDB可用）
SELECT create_hypertable('system_logs', 'timestamp', if_not_exists => TRUE);

-- 创建索引
CREATE INDEX IF NOT EXISTS idx_market_data_symbol_timestamp ON market_data(symbol, timestamp DESC);
CREATE INDEX IF NOT EXISTS idx_orders_strategy_id ON orders(strategy_id);
CREATE INDEX IF NOT EXISTS idx_orders_symbol ON orders(symbol);
CREATE INDEX IF NOT EXISTS idx_orders_status ON orders(status);
CREATE INDEX IF NOT EXISTS idx_positions_strategy_id ON positions(strategy_id);
CREATE INDEX IF NOT EXISTS idx_trades_strategy_id ON trades(strategy_id);
CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol);
CREATE INDEX IF NOT EXISTS idx_risk_metrics_strategy_id ON risk_metrics(strategy_id);
CREATE INDEX IF NOT EXISTS idx_system_logs_level ON system_logs(level);
CREATE INDEX IF NOT EXISTS idx_system_logs_timestamp ON system_logs(timestamp DESC);

-- 创建更新时间触发器函数
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ language 'plpgsql';

-- 为需要的表创建更新时间触发器
CREATE TRIGGER update_users_updated_at BEFORE UPDATE ON users
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_strategies_updated_at BEFORE UPDATE ON strategies
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_orders_updated_at BEFORE UPDATE ON orders
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
EOF

    # 执行SQL脚本
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -f /tmp/create_tables.sql
    
    if [ $? -eq 0 ]; then
        log_success "数据表结构创建成功"
    else
        log_error "数据表结构创建失败"
        return 1
    fi
    
    # 清理临时文件
    rm -f /tmp/create_tables.sql
}

# 插入初始数据
insert_initial_data() {
    log_info "插入初始数据..."
    
    # 创建初始数据SQL脚本
    cat > /tmp/initial_data.sql << 'EOF'
-- 插入默认管理员用户
INSERT INTO users (username, email, password_hash, is_admin) 
VALUES ('admin', 'admin@aiquant.com', '$2b$12$LQv3c1yqBWVHxkd0LHAkCOYz6TtxMQJqhN8/LewdBPj3QJgusgqHu', true)
ON CONFLICT (username) DO NOTHING;

-- 插入示例策略
INSERT INTO strategies (name, description, strategy_type, parameters) 
VALUES 
    ('均线策略', '基于移动平均线的交易策略', 'trend_following', '{"short_period": 5, "long_period": 20}'),
    ('RSI策略', '基于相对强弱指数的交易策略', 'mean_reversion', '{"period": 14, "oversold": 30, "overbought": 70}'),
    ('布林带策略', '基于布林带的交易策略', 'volatility', '{"period": 20, "std_dev": 2}')
ON CONFLICT DO NOTHING;

-- 插入示例风险指标阈值
INSERT INTO risk_metrics (strategy_id, metric_name, threshold_value, status)
SELECT 
    s.id,
    metric_name,
    threshold_value,
    'normal'
FROM strategies s
CROSS JOIN (
    VALUES 
        ('max_drawdown', 0.05),
        ('var_95', 0.02),
        ('sharpe_ratio', 1.0),
        ('position_concentration', 0.1)
) AS metrics(metric_name, threshold_value)
ON CONFLICT DO NOTHING;
EOF

    # 执行初始数据SQL脚本
    PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -f /tmp/initial_data.sql
    
    if [ $? -eq 0 ]; then
        log_success "初始数据插入成功"
    else
        log_warning "初始数据插入可能有问题，请检查"
    fi
    
    # 清理临时文件
    rm -f /tmp/initial_data.sql
}

# 验证数据库
verify_database() {
    log_info "验证数据库结构..."
    
    # 检查表是否创建成功
    local tables=$(PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -t -c "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';")
    
    if [ "$tables" -gt 0 ]; then
        log_success "数据库验证成功，共创建 $tables 个表"
        
        # 显示表列表
        log_info "数据表列表:"
        PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME -c "\dt"
    else
        log_error "数据库验证失败，未找到数据表"
        return 1
    fi
}

# 备份数据库结构
backup_schema() {
    log_info "备份数据库结构..."
    
    local backup_file="backups/schema_$(date '+%Y%m%d_%H%M%S').sql"
    
    if [ ! -d "backups" ]; then
        mkdir -p backups
    fi
    
    PGPASSWORD=$DB_PASSWORD pg_dump -h $DB_HOST -p $DB_PORT -U $DB_USER -d $DB_NAME --schema-only > "$backup_file"
    
    if [ $? -eq 0 ]; then
        log_success "数据库结构备份完成: $backup_file"
    else
        log_warning "数据库结构备份失败"
    fi
}

# 显示连接信息
show_connection_info() {
    echo ""
    echo "数据库连接信息:"
    echo "  主机: $DB_HOST"
    echo "  端口: $DB_PORT"
    echo "  数据库: $DB_NAME"
    echo "  用户: $DB_USER"
    echo ""
    echo "连接字符串:"
    echo "  postgresql://$DB_USER:****@$DB_HOST:$DB_PORT/$DB_NAME"
}

# 主函数
main() {
    log_info "开始初始化AIQuant数据库..."
    
    read_config
    
    case "${1:-}" in
        --help|-h)
            echo "AIQuant数据库初始化脚本"
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --help, -h        显示帮助信息"
            echo "  --drop            删除现有数据库后重新创建"
            echo "  --schema-only     仅创建表结构，不插入初始数据"
            echo "  --verify          仅验证数据库结构"
            echo "  --backup          备份数据库结构"
            echo ""
            echo "环境变量:"
            echo "  DATABASE_HOST     数据库主机 (默认: localhost)"
            echo "  DATABASE_PORT     数据库端口 (默认: 5432)"
            echo "  DATABASE_NAME     数据库名称 (默认: aiquant)"
            echo "  DATABASE_USER     数据库用户 (默认: aiquant)"
            echo "  DATABASE_PASSWORD 数据库密码 (默认: aiquant123)"
            exit 0
            ;;
        --drop)
            log_warning "删除现有数据库..."
            PGPASSWORD=$DB_PASSWORD psql -h $DB_HOST -p $DB_PORT -U $DB_USER -d postgres -c "DROP DATABASE IF EXISTS $DB_NAME;"
            ;;
        --verify)
            check_postgres && verify_database
            exit 0
            ;;
        --backup)
            check_postgres && backup_schema
            exit 0
            ;;
        --schema-only)
            check_postgres || exit 1
            create_database || exit 1
            install_extensions || exit 1
            create_tables || exit 1
            verify_database || exit 1
            backup_schema
            show_connection_info
            log_success "数据库结构初始化完成!"
            exit 0
            ;;
    esac
    
    # 完整初始化流程
    check_postgres || exit 1
    create_database || exit 1
    install_extensions || exit 1
    create_tables || exit 1
    insert_initial_data || exit 1
    verify_database || exit 1
    backup_schema
    
    show_connection_info
    log_success "AIQuant数据库初始化完成!"
}

# 执行主函数
main "$@"